{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09f1b46",
   "metadata": {},
   "outputs": [],
   "source": [
    "#1.1\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "class army_group(object):\n",
    "    def __init__(self,bird_num=100,eagle_num=100,para=[2,10,3,1,1,5],e_para=[2,10,3,1,1,5],center=[50,50]):#\n",
    "        self.bird_num=bird_num\n",
    "        self.eagle_num=eagle_num\n",
    "        self.bird_map=np.random.rand(bird_num,2)\n",
    "        self.eagle_map=np.random.rand(eagle_num,2)\n",
    "        self.vmap=np.random.randn(bird_num,2)\n",
    "        self.e_vmap=np.random.randn(eagle_num,2)\n",
    "        self.amap=np.zeros([bird_num,2])\n",
    "        self.e_amap=np.zeros([eagle_num,2])\n",
    "        self.dt=0.1\n",
    "        self.para=para\n",
    "        self.e_para=e_para\n",
    "        self.alim=10\n",
    "        self.center=center\n",
    "        self.vmax=5\n",
    "    def ini_2group(self):\n",
    "        for i in range(int(self.bird_num)):\n",
    "            self.bird_map[i,0]=i*(100/self.bird_num)\n",
    "            self.bird_map[i,1]=10+np.random.rand()\n",
    "            self.vmap[i]=np.array([0,5])\n",
    "        for i in range(self.eagle_num):\n",
    "            self.eagle_map[i,0]=(i)*(100/self.eagle_num)\n",
    "            self.eagle_map[i,1]=90+np.random.rand()\n",
    "            self.e_vmap[i]=np.array([0,-5])\n",
    "    def ini_2group_ran(self):\n",
    "        for i in range(int(self.bird_num)):\n",
    "            self.bird_map[i,0]=np.random.rand()*100\n",
    "            self.bird_map[i,1]=np.random.rand()*100\n",
    "            self.vmap[i]=np.array([0,0])\n",
    "        for i in range(self.eagle_num):\n",
    "            self.eagle_map[i,0]=50#np.random.rand()*100\n",
    "            self.eagle_map[i,1]=50#np.random.rand()*100\n",
    "            self.e_vmap[i]=np.array([0,0])\n",
    "    \n",
    "    def main(self,end_num=100,show=True,check_point=1):\n",
    "        check_line=[]\n",
    "        self.ini_2group()\n",
    "        for epoch in range(end_num):\n",
    "            self.step()\n",
    "            \n",
    "            if epoch%check_point==0:\n",
    "                if show:\n",
    "                    self.show_plt(epoch)\n",
    "                check_line.append(self.check())\n",
    "            if len(self.bird_map)==0 or len(self.eagle_map)==0:\n",
    "                break\n",
    "        return np.array(check_line)\n",
    "    def meet(self,):\n",
    "        check_line=self.main(check_point=1)\n",
    "        return check_line\n",
    "    def show_plt(self,t):\n",
    "        fig=plt.figure(figsize=(10,10))\n",
    "        plt.xlim(0,100)\n",
    "        plt.ylim(0,100)\n",
    "        plt.scatter(self.center[0], self.center[1], c='yellow', s=100, marker='*')\n",
    "        \n",
    "        for i in range(len(self.bird_map)):\n",
    "            plt.scatter(self.bird_map[i,0],self.bird_map[i,1],color=\"blue\")\n",
    "            plt.annotate(\"\",(self.bird_map[i,0]+self.vmap[i,0]*self.dt,self.bird_map[i,1]+self.vmap[i,1]*self.dt),(self.bird_map[i,0],self.bird_map[i,1]),arrowprops=dict(arrowstyle=\"->\"))\n",
    "        for i in range(len(self.eagle_map)):\n",
    "            plt.scatter(self.eagle_map[i,0],self.eagle_map[i,1],color=\"red\")\n",
    "            plt.annotate(\"\",(self.eagle_map[i,0]+self.e_vmap[i,0]*self.dt,self.eagle_map[i,1]+self.e_vmap[i,1]*self.dt),(self.eagle_map[i,0],self.eagle_map[i,1]),arrowprops=dict(arrowstyle=\"->\"))\n",
    "        plt.savefig(SAVE_PATH+str(t)+\".jpg\")\n",
    "        plt.show()\n",
    "        return None\n",
    "    def distance(self,a,b):\n",
    "        return ((a[0]-b[0])**2+(a[1]-b[1])**2)**0.5\n",
    "    def fight(self,group_0,group_1,v_group_0,v_group_1,a_group_0,a_group_1,para_0,para_1,flee=5):\n",
    "        del_num=0\n",
    "        temp_vmap,temp_bird_map,temp_amap=v_group_0,group_0,a_group_0\n",
    "        for i in range(len(group_0)):\n",
    "            al_num=0\n",
    "            foe_num=0\n",
    "            al_num_w=0\n",
    "            foe_num_w=0\n",
    "            nearest_foe=self.center\n",
    "            nearest_d=None\n",
    "            for m in range(len(group_0)):\n",
    "                d_temp=self.distance(group_0[i],group_0[m])\n",
    "                if d_temp<para_0[3]:\n",
    "                    al_num+=para_0[4]\n",
    "                if d_temp<para_0[5]:\n",
    "                    al_num_w+=para_0[4]\n",
    "                \n",
    "            for m in range(len(group_1)):\n",
    "                d_temp=self.distance(group_0[i],group_1[m])\n",
    "                if nearest_d==None or nearest_d>d_temp:\n",
    "                    nearest_d=d_temp\n",
    "                    nearest_foe=group_1[m]\n",
    "                if d_temp<para_1[3]:\n",
    "                    foe_num+=para_1[4]\n",
    "                if d_temp<para_1[5]:\n",
    "                    foe_num_w+=para_1[4]\n",
    "                \n",
    "            if foe_num>al_num:#敌方占优时死亡\n",
    "                temp_bird_map=np.delete(temp_bird_map,i-del_num,axis=0)\n",
    "                temp_vmap=np.delete(temp_vmap,i-del_num,axis=0)\n",
    "                temp_amap=np.delete(temp_amap,i-del_num,axis=0)\n",
    "                del_num+=1\n",
    "            elif foe_num>0:#攻击时立定 \n",
    "                v_group_0[i,0],v_group_0[i,1]=0,0      \n",
    "            elif foe_num_w>al_num_w:#若视野内敌方占优，则撤退\n",
    "                temp_amap[i-del_num,0]-=flee*para_0[0]*(nearest_foe[0]-group_0[i,0])\n",
    "                temp_amap[i-del_num,1]-=flee*para_0[0]*(nearest_foe[1]-group_0[i,1])\n",
    "            else:#elif foe_num_w>0:#若视野内自方占优，则进攻\n",
    "                temp_amap[i-del_num,0]+=para_0[0]*(nearest_foe[0]-group_0[i,0])#向心性\n",
    "                temp_amap[i-del_num,1]+=para_0[0]*(nearest_foe[1]-group_0[i,1])\n",
    "        e_temp_vmap,e_temp_bird_map,e_temp_amap=v_group_1,group_1,a_group_1\n",
    "        del_num=0\n",
    "        for i in range(len(group_1)):\n",
    "            al_num=0\n",
    "            foe_num=0\n",
    "            al_num_w=0\n",
    "            foe_num_w=0\n",
    "            nearest_foe=self.center\n",
    "            nearest_d=None\n",
    "            for m in range(len(group_1)):\n",
    "                d_temp=self.distance(group_1[i],group_1[m])\n",
    "                if d_temp<para_1[3]:\n",
    "                    al_num+=para_1[4]\n",
    "                if d_temp<para_1[5]:\n",
    "                    al_num_w+=para_1[4]\n",
    "                \n",
    "            for m in range(len(group_0)):\n",
    "                d_temp=self.distance(group_1[i],group_0[m])\n",
    "                if nearest_d==None or nearest_d>d_temp:\n",
    "                    nearest_d=d_temp\n",
    "                    nearest_foe=group_0[m]\n",
    "                if d_temp<para_0[3]:\n",
    "                    foe_num+=para_0[4]\n",
    "                if d_temp<para_0[5]:\n",
    "                    foe_num_w+=para_0[4]\n",
    "                \n",
    "            if foe_num>al_num:#敌方占优时死亡\n",
    "                e_temp_bird_map=np.delete(e_temp_bird_map,i-del_num,axis=0)\n",
    "                e_temp_vmap=np.delete(e_temp_vmap,i-del_num,axis=0)\n",
    "                e_temp_amap=np.delete(e_temp_amap,i-del_num,axis=0)\n",
    "                del_num+=1\n",
    "            elif foe_num>0:#攻击时立定 \n",
    "                v_group_1[i,0],v_group_1[i,1]=0,0      \n",
    "            elif foe_num_w>al_num_w:\n",
    "                e_temp_amap[i-del_num,0]-=flee*para_1[0]*(nearest_foe[0]-group_1[i,0])#向心性\n",
    "                e_temp_amap[i-del_num,1]-=flee*para_1[0]*(nearest_foe[1]-group_1[i,1])\n",
    "            else:#elif foe_num_w>0:\n",
    "                e_temp_amap[i-del_num,0]+=para_1[0]*(nearest_foe[0]-group_1[i,0])#向心性\n",
    "                e_temp_amap[i-del_num,1]+=para_1[0]*(nearest_foe[1]-group_1[i,1])\n",
    "             \n",
    "        v_group_0,group_0,a_group_0=temp_vmap,temp_bird_map,temp_amap\n",
    "        v_group_1,group_1,a_group_1=e_temp_vmap,e_temp_bird_map,e_temp_amap\n",
    "        return temp_vmap,temp_bird_map,temp_amap,e_temp_vmap,e_temp_bird_map,e_temp_amap\n",
    "    def step(self,to_center=0.5):\n",
    "        self.bird_map=self.bird_map+self.vmap*self.dt\n",
    "        self.eagle_map=self.eagle_map+self.e_vmap*self.dt\n",
    "        self.amap=np.zeros([len(self.bird_map),2])\n",
    "        self.e_amap=np.zeros([len(self.eagle_map),2])\n",
    "        #(self,group_0,group_1,v_group_0,v_group_1,a_group_0,a_group_1,para_0,para_1)\n",
    "        self.vmap,self.bird_map,self.amap,self.e_vmap,self.eagle_map,self.e_amap=self.fight(self.bird_map,self.eagle_map,self.vmap,self.e_vmap,self.amap,self.e_amap,self.para,self.e_para)\n",
    "        #temp_vmap,temp_bird_map,temp_amap,e_temp_vmap,e_temp_bird_map,e_temp_amap\n",
    "        for i in range(len(self.bird_map)):\n",
    "            near_g,near_v=self.near_group(self_num=i)\n",
    "            #self.amap[i,0]+=to_center*self.para[0]*(self.center[0]-self.bird_map[i,0])#争夺center\n",
    "            #self.amap[i,1]+=to_center*self.para[0]*(self.center[1]-self.bird_map[i,1])\n",
    "            self.amap[i,0]+=self.para[0]*(np.average(near_g[:,0])-self.bird_map[i,0])#向心性\n",
    "            self.amap[i,1]+=self.para[0]*(np.average(near_g[:,1])-self.bird_map[i,1])\n",
    "            self.amap[i,0]+=self.para[1]*(np.average(near_v[:,0])-self.vmap[i,0])#同速性\n",
    "            self.amap[i,1]+=self.para[1]*(np.average(near_v[:,1])-self.vmap[i,1])\n",
    "        \n",
    "            for j in range(len(near_g)):\n",
    "                if near_g[j,0]-self.bird_map[i,0]>0:\n",
    "                    x0=-1\n",
    "                else:\n",
    "                    x0=1\n",
    "                if near_g[j,1]-self.bird_map[i,1]>0:\n",
    "                    x1=-1\n",
    "                else:\n",
    "                    x1=1\n",
    "                if near_g[j,0]-self.bird_map[i,0]!=0 and near_g[j,1]-self.bird_map[i,1]!=0:\n",
    "                    self.amap[i,0]+=x0*self.para[2]/(near_g[j,0]-self.bird_map[i,0])**2#\n",
    "                    self.amap[i,1]+=x1*self.para[2]/(near_g[j,1]-self.bird_map[i,1])**2\n",
    "            if self.amap[i,0]>self.alim:\n",
    "                self.amap[i,0]=self.alim\n",
    "            if self.amap[i,0]<-self.alim:\n",
    "                self.amap[i,0]=-self.alim\n",
    "            if self.amap[i,1]>self.alim:\n",
    "                self.amap[i,1]=self.alim\n",
    "            if self.amap[i,1]<-self.alim:\n",
    "                self.amap[i,1]=-self.alim            \n",
    "        self.vmap=self.vmap+self.amap*self.dt\n",
    "        for i in range(len(self.vmap)):\n",
    "            if self.vmap[i,0]**2+self.vmap[i,1]**2>self.vmax**2:\n",
    "                self.vmap[i]=self.vmap[i]*self.vmax/(self.vmap[i,0]**2+self.vmap[i,1]**2)**0.5\n",
    "        for i in range(len(self.eagle_map)):\n",
    "            e_near_g,e_near_v=self.near_group_e(self_num=i)\n",
    "            self.e_amap[i,0]+=self.para[0]*(self.center[0]-self.eagle_map[i,0])\n",
    "            self.e_amap[i,1]+=self.para[0]*(self.center[1]-self.eagle_map[i,1])\n",
    "            self.e_amap[i,0]+=self.e_para[0]*(np.average(e_near_g[:,0])-self.eagle_map[i,0])\n",
    "            self.e_amap[i,1]+=self.e_para[0]*(np.average(e_near_g[:,1])-self.eagle_map[i,1])\n",
    "            self.e_amap[i,0]+=self.e_para[1]*(np.average(e_near_v[:,0])-self.e_vmap[i,0])\n",
    "            self.e_amap[i,1]+=self.e_para[1]*(np.average(e_near_v[:,1])-self.e_vmap[i,1])\n",
    "            for j in range(len(e_near_g)):\n",
    "                if e_near_g[j,0]-self.eagle_map[i,0]>0:\n",
    "                    x0=-1\n",
    "                else:\n",
    "                    x0=1\n",
    "                if e_near_g[j,1]-self.eagle_map[i,1]>0:\n",
    "                    x1=-1\n",
    "                else:\n",
    "                    x1=1\n",
    "                if e_near_g[j,0]-self.eagle_map[i,0]!=0 and e_near_g[j,1]-self.eagle_map[i,1]!=0:\n",
    "                    self.e_amap[i,0]+=x0*self.e_para[2]/(e_near_g[j,0]-self.eagle_map[i,0])**2\n",
    "                    self.e_amap[i,1]+=x1*self.e_para[2]/(e_near_g[j,1]-self.eagle_map[i,1])**2\n",
    "            if self.e_amap[i,0]>self.alim:\n",
    "                self.e_amap[i,0]=self.alim\n",
    "            if self.e_amap[i,0]<-self.alim:\n",
    "                self.e_amap[i,0]=-self.alim\n",
    "            if self.e_amap[i,1]>self.alim:\n",
    "                self.e_amap[i,1]=self.alim\n",
    "            if self.e_amap[i,1]<-self.alim:\n",
    "                self.e_amap[i,1]=-self.alim  \n",
    "        self.e_vmap=self.e_vmap+self.e_amap*self.dt\n",
    "        for i in range(len(self.e_vmap)):\n",
    "            if self.e_vmap[i,0]**2+self.e_vmap[i,1]**2>self.vmax**2:\n",
    "                self.e_vmap[i]=self.e_vmap[i]*self.vmax/(self.e_vmap[i,0]**2+self.e_vmap[i,1]**2)**0.5\n",
    "        return None\n",
    "    def near_group(self,self_num=0,d=5):\n",
    "        output=[]\n",
    "        voutput=[]\n",
    "        for i in range(len(self.bird_map)):\n",
    "            if self.bird_map[i,0]-self.bird_map[self_num,0]<d and self.bird_map[i,1]-self.bird_map[self_num,1]<d:\n",
    "                output.append(self.bird_map[i])\n",
    "                voutput.append(self.vmap[i])\n",
    "        return np.array(output),np.array(voutput)\n",
    "    def near_group_e(self,self_num=0,d=5):\n",
    "        output=[]\n",
    "        voutput=[]\n",
    "        for i in range(len(self.eagle_map)):\n",
    "            if self.eagle_map[i,0]-self.eagle_map[self_num,0]<d and self.eagle_map[i,1]-self.eagle_map[self_num,1]<d:\n",
    "                output.append(self.eagle_map[i])\n",
    "                voutput.append(self.e_vmap[i])\n",
    "        return np.array(output),np.array(voutput)\n",
    "        return np.array(output),np.array(voutput)\n",
    "    def check(self):\n",
    "        return [len(self.bird_map),len(self.eagle_map)]\n",
    "\n",
    "SAVE_PATH=\"D:/3_bodies/DLA_pic/fight_\"  #改为合适的目录\n",
    "army_0=army_group(bird_num=200,eagle_num=200,para=[2,10,1,2,1,5],e_para=[2,10,1,2,1,5])\n",
    "a_line=army_0.main(end_num=10000)\n",
    "plt.plot(a_line[:][0])\n",
    "plt.plot(a_line[:][1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d408d6e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#imgs2video\n",
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "fourcc = cv2.VideoWriter_fourcc('X', 'V', 'I', 'D') \n",
    "path=r\"D:/3_bodies/DLA_pic/fight_\"#input(\"path?\")\n",
    "save_path=r\"D:/3_bodies/fight_2-2_r_0.mp4\"#input(\"save_path?\")\n",
    "post_video=cv2.VideoWriter(save_path, fourcc, 30,(720,720))\n",
    "len_img=4428\n",
    "for i in range(len_img+1):\n",
    "    print(path+str(i)+\".jpg\")\n",
    "    img=cv2.imread(path+str(i)+\".jpg\")\n",
    "    post_video.write(img)\n",
    "    if i==len_img:\n",
    "        for t in range(100):\n",
    "            post_video.write(img)\n",
    "    del img\n",
    "    #if filename==\"0.jpg\":\n",
    "        #Image.fromarray(img).show()\n",
    "\n",
    "post_video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d221447",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Superman version\n",
    "army_0=army_group(bird_num=1000,eagle_num=1,para=[2,10,3,2,1,5],e_para=[2,10,3,2,10,5])\n",
    "army_0.main(end_num=10000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65db345b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "def lancaster_func(x_0,y_0,a=1,b=1,dt=0.01,plt_show=True):\n",
    "    \n",
    "    x_line,y_line=[x_0],[y_0]\n",
    "    while True:\n",
    "        x_line.append(x_line[-1]-a*y_line[-1]*dt)\n",
    "        y_line.append(y_line[-1]-b*x_line[-1]*dt)\n",
    "        if x_line[-1]<1 or y_line[-1]<1:\n",
    "            break\n",
    "    if plt_show:\n",
    "        plt.plot(x_line)\n",
    "        plt.plot(y_line)\n",
    "        plt.show()\n",
    "    return x_line,y_line\n",
    "lancaster_func(600,610)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b56f6861",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(a_line[:,0])\n",
    "plt.plot(a_line[:,1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c15157",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
